package com.pai4j.aigc.llm;
import com.fasterxml.jackson.core.type.TypeReference;
import com.pai4j.common.bean.PAIResponseBean;
import com.pai4j.common.bean.PAIResponseBeanUtil;
import com.pai4j.common.constants.MilvusConstants;
import com.pai4j.common.enums.PromptTemplateEnum;
import com.pai4j.common.service.FreeMarkEngineService;
import com.pai4j.common.util.JsonUtil;
import com.pai4j.domain.vo.llm.ChatCompletionMessage;
import com.pai4j.domain.dto.llm.ChatMessageRole;
import com.pai4j.domain.vo.request.request.mess.ChatMessageRequestVO;
import com.pai4j.domain.vo.request.search.MilvusSearchRequestVO;
import com.pai4j.domain.vo.response.search.ArticleMilvusSearchResponse;
import com.pai4j.remote.search.VectorSearchEngineClient;
import com.pai4j.aigc.llm.entity.LLMModelEntity;
import com.pai4j.aigc.llm.service.LLMModelService;
import com.pai4j.aigc.llm.service.LLMUsageService;
import com.pai4j.aigc.llm.service.MetricsService;
import com.pai4j.aigc.llm.service.PricingService;
import com.pai4j.aigc.llm.service.TokenEstimator;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.*;


@RefreshScope
@Service
public class AiBotQAService {


    @Autowired
    private VectorSearchEngineClient vectorSearchEngineClient;

    @Autowired
    private FreeMarkEngineService freeMarkEngineService;

    @Value("${llm.config.model}")
    private String model;

    @Autowired
    private LLMUsageService usageService;

    @Autowired
    private TokenEstimator tokenEstimator;

    @Autowired
    private PricingService pricingService;

    @Autowired
    private MetricsService metricsService;

    @Autowired
    private LLMModelService modelService;

    /**
     * 向量相似度搜索默认取相似度最高前两条数据
     */
    private static final int SEARCH_TOP_K = 2;

    /**
     * 相似度限制条件
     */
    private static final float MIN_SCORE = 0.4f;

    private static final String AIBOT_ROLE = "你是Chat4j项目的Ai助手，任务是根据用户给出的参考数据，基于输入query，精准的给出答案。";

    /**
     * 我的ai助手问答场景实现 （RAG）
     *
     * @param sseEmitter
     * @param chatMessage
     * @return
     */
    public String chat(SseEmitter sseEmitter, ChatMessageRequestVO chatMessage) {

        String query = chatMessage.getMessage();
        String senderId = chatMessage.getSenderId();
        /**
         * 基于公众号文章相似度实现搜索
         */
        List<ArticleMilvusSearchResponse> searchResponses = this.search(senderId, query);

        /**
         * prompt工程构件需要大模型处理指令
         */
        Map<String, Object> vars = new HashMap<>();
        vars.put("searchResponses", searchResponses);
        vars.put("input", query);
        String prompt;
        try {
            prompt = freeMarkEngineService.getContentByTemplate(PromptTemplateEnum.AIBOT_RAG.getPath(), vars);
        } catch (Exception e) {
            this.sseSend(sseEmitter, "prompt构建失败");
            e.printStackTrace();
            return null;
        }
        /**
         * 增强
         */
        // 使用工厂模式，取对应的LLM服务，来完成对话
        AbstractLLMChatService abstractLLMChatService = LLMServiceFactory.getLLMService(model);
        final List<ChatCompletionMessage> messageList = Arrays.asList(
                new ChatCompletionMessage(ChatMessageRole.SYSTEM.value(), AIBOT_ROLE),
                new ChatCompletionMessage(ChatMessageRole.USER.value(), prompt));
        String aiResponse;
        // resolve model info
        java.util.Optional<LLMModelEntity> modelEntityOpt = modelService.findByGlobalCode(model);
        LLMModelEntity modelEntity = modelEntityOpt.orElse(null);
        String providerKey = modelEntity != null ? modelEntity.getProvider() : model;
        Long modelId = modelEntity != null ? modelEntity.getId() : -1L;
        String requestId = java.util.UUID.randomUUID().toString();
        long start = System.currentTimeMillis();
        try {
            aiResponse = abstractLLMChatService.chat(sseEmitter, messageList);
            long latency = System.currentTimeMillis() - start;
            int pt = tokenEstimator.estimateTokens(prompt);
            int ct = tokenEstimator.estimateTokens(aiResponse);
            int total = pt + ct;
            Integer costCents = pricingService.calcCostCents(modelEntity, pt, ct);
            usageService.recordSuccess(null, requestId, chatMessage.getSenderId(), modelId, providerKey,
                    modelEntity != null ? modelEntity.getCode() : providerKey,
                    pt, ct, total, latency, costCents);
            metricsService.recordSuccess(providerKey, null, modelEntity != null ? modelEntity.getCode() : providerKey, latency);
        } catch (Exception e) {
            long latency = System.currentTimeMillis() - start;
            usageService.recordFailure(null, requestId, chatMessage.getSenderId(), modelId, providerKey,
                    modelEntity != null ? modelEntity.getCode() : providerKey,
                    "LLM_ERROR", e.getMessage(), latency);
            metricsService.recordFailure(providerKey, null, modelEntity != null ? modelEntity.getCode() : providerKey, latency);
            this.sseSend(sseEmitter, "服务器繁忙");
            return null;
        }
        return aiResponse;
    }

    /**
     * RPC 远程请求我们的搜索引擎服务，基于向量相似度检索，搜索topK条公众号文章片段
     *
     * @param userId
     * @param query
     * @return
     */
    private List<ArticleMilvusSearchResponse> search(String userId, String query) {
        String biz = MilvusConstants.SearchBiz.SEARCH_ARTICLE.getBiz();
        MilvusSearchRequestVO req = MilvusSearchRequestVO.builder()
                .userId(userId)
                .query(query)
                .score(MIN_SCORE)
                .topK(SEARCH_TOP_K)
                .build();
        PAIResponseBean<Object> response = vectorSearchEngineClient.search(biz, req);
        if (PAIResponseBeanUtil.isOk(response)) {
            return JsonUtil.fromJson(JsonUtil.toJsonString(response.getData()), new TypeReference<List<ArticleMilvusSearchResponse>>() { });
        }
        return Collections.emptyList();
    }

    /**
     * sse 消息推送
     *
     * @param message
     */
    public static void sseSend(SseEmitter sseEmitter, String message) {
        try {
            sseEmitter.send(message);
            sseEmitter.send("finished");
        } catch (IOException ex) {
            throw new RuntimeException(ex);
        }
    }
}
